{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cb0aacbb",
   "metadata": {
    "cell_id": "c4ca66ed8748422fa04e280de76ec7bc",
    "colab_type": "text",
    "deepnote_cell_type": "markdown",
    "id": "TKvYiJgnYExi"
   },
   "source": [
    "# Instructions\n",
    "Run the first couple of cells.\n",
    "\n",
    "When you get to the generated Meshcat URL, click to open it. Read the notes below the cell to set up your gamepad or see/set the velocity command values mapped to your keyboard.\n",
    "**note:** driving with the gamepad is a much better experience!!\n",
    "\n",
    "Then, run the final cell and drive the robot around! Observe the many realistic forces acting upon the robot (contact forces, gravity, inertia, etc). The only way that the robot is moving is because the motors apply a torque at the wheel and Drake simulates the detailed physics of the mecanum wheel rollers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa147216",
   "metadata": {
    "cell_id": "bf38e76b2fe94ba2aa0021a8a2b4b6aa",
    "colab": {},
    "colab_type": "code",
    "deepnote_cell_type": "code",
    "deepnote_to_be_reexecuted": false,
    "execution_millis": 722,
    "execution_start": 1701625998486,
    "id": "A4QOaw_zYLfI",
    "lines_to_next_cell": 2,
    "source_hash": "c3b5808c"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    DiagramBuilder,\n",
    "    DiscreteContactSolver,\n",
    "    LeafSystem,\n",
    "    MeshcatVisualizer,\n",
    "    Parser,\n",
    "    Simulator,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.meshcat_utils import MeshcatSliders\n",
    "from manipulation.utils import ConfigureParser, FindResource"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e94fb979",
   "metadata": {
    "cell_id": "b43a3e413013413bb64027d0f54c390b",
    "deepnote_cell_type": "code",
    "deepnote_to_be_reexecuted": false,
    "execution_millis": 168,
    "execution_start": 1701625999210,
    "source_hash": "6d42057f"
   },
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e4d2adf",
   "metadata": {
    "cell_id": "525e0f538345457fae0e4f4d37b188ef",
    "deepnote_cell_type": "text-cell-p",
    "formattedRanges": [
     {
      "fromCodePoint": 0,
      "marks": {
       "bold": true
      },
      "toCodePoint": 95,
      "type": "marks"
     }
    ],
    "is_collapsed": false
   },
   "source": [
    "If you want to use a gamepad to steer, press some buttons on the gamepad now with the meshcat window in focus."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c86b8602",
   "metadata": {
    "cell_id": "951104aa0d42481b9eb40c772a89004e",
    "colab_type": "text",
    "deepnote_cell_type": "markdown",
    "id": "4cTkwpJU8tGX"
   },
   "source": [
    "You can control the vehicle with either a gamepad (if you have it) or with a keyboard teleop (if you don't). If you're using the keyboard, then click \"Open Controls\" in meshcat to see/set the velocity command values. The sliders are also mapped to the keys:\n",
    "```\n",
    "xdot : [ArrowDown, ArrowUp]\n",
    "ydot : [ArrowLeft, ArrowRight]\n",
    "yawdot : [KeyA, KeyD]\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdade946",
   "metadata": {
    "cell_id": "30816fe345074130b9e1a52f4a0ea0f7",
    "colab": {},
    "colab_type": "code",
    "deepnote_cell_type": "code",
    "deepnote_to_be_reexecuted": false,
    "execution_millis": 677737,
    "execution_start": 1701625999382,
    "id": "8wj7ZlyEw_AB",
    "source_hash": "343970e2"
   },
   "outputs": [],
   "source": [
    "class GamepadCommand(LeafSystem):\n",
    "    def __init__(self, meshcat):\n",
    "        LeafSystem.__init__(self)\n",
    "\n",
    "        self._meshcat = meshcat\n",
    "        port = self.DeclareVectorOutputPort(\"command\", 3, self.CalcOutput)\n",
    "        port.disable_caching_by_default()\n",
    "\n",
    "    def CalcOutput(self, context, output):\n",
    "        gamepad = self._meshcat.GetGamepad()\n",
    "        assert gamepad.index != None\n",
    "\n",
    "        # https://beej.us/blog/data/javascript-gamepad/\n",
    "        def CreateStickDeadzone(x, y):\n",
    "            stick = np.array([x, y])\n",
    "            deadzone = 0.2\n",
    "            m = np.linalg.norm(stick)\n",
    "            if m < deadzone:\n",
    "                return np.array([0, 0])\n",
    "            over = (m - deadzone) / (1 - deadzone)\n",
    "            return stick * over / m\n",
    "\n",
    "        left = CreateStickDeadzone(gamepad.axes[0], gamepad.axes[1])\n",
    "        right = CreateStickDeadzone(gamepad.axes[2], gamepad.axes[3])\n",
    "\n",
    "        output[0] = -left[1]  # Left stick y => vx\n",
    "        output[1] = -left[0]  # Left stick x => vy\n",
    "        output[2] = right[0]  # Right stick x => wz\n",
    "\n",
    "\n",
    "class MyController(LeafSystem):\n",
    "    def __init__(self, plant, model_instance):\n",
    "        LeafSystem.__init__(self)\n",
    "        self._wheel_velocity_indices = (\n",
    "            np.array(\n",
    "                [\n",
    "                    plant.GetJointByName(\n",
    "                        \"front_left_wheel_joint\", model_instance\n",
    "                    ).velocity_start(),\n",
    "                    plant.GetJointByName(\n",
    "                        \"front_right_wheel_joint\", model_instance\n",
    "                    ).velocity_start(),\n",
    "                    plant.GetJointByName(\n",
    "                        \"rear_left_wheel_joint\", model_instance\n",
    "                    ).velocity_start(),\n",
    "                    plant.GetJointByName(\n",
    "                        \"rear_right_wheel_joint\", model_instance\n",
    "                    ).velocity_start(),\n",
    "                ]\n",
    "            )\n",
    "            + plant.num_positions()\n",
    "        )\n",
    "\n",
    "        # command is the [vx, vy, wz] components of V_WRobot_Robot.\n",
    "        self.DeclareVectorInputPort(\"command\", 3)\n",
    "        self.DeclareVectorInputPort(\"state\", plant.num_multibody_states())\n",
    "        self.DeclareVectorOutputPort(\"motor_torque\", 4, self.CalcTorques)\n",
    "\n",
    "        # These should match the parameters used to create the URDF.\n",
    "        wheel_radius = 0.045 + (\n",
    "            0.015 / 2\n",
    "        )  # hub_radius + (roller_diameter / 2).\n",
    "        wheelbase = 0.2\n",
    "        track = 0.205\n",
    "        self._wheel_velocity_kp = 0.1\n",
    "\n",
    "        lx = track * 0.5\n",
    "        ly = wheelbase * 0.5\n",
    "\n",
    "        # From anzu/punito/control/mecanum_kinematics.cc\n",
    "        self._vehicle_to_wheel_map = (\n",
    "            np.array(\n",
    "                [\n",
    "                    [1, -1, (lx + ly)],\n",
    "                    [1, 1, -(lx + ly)],\n",
    "                    [1, 1, (lx + ly)],\n",
    "                    [1, -1, -(lx + ly)],\n",
    "                ]\n",
    "            )\n",
    "            / wheel_radius\n",
    "        )\n",
    "\n",
    "    def CalcTorques(self, context, output):\n",
    "        # From anzu/punito/sim/robot_master_controller.cc\n",
    "        command = self.get_input_port(0).Eval(context)\n",
    "        state = self.get_input_port(1).Eval(context)\n",
    "        wheel_velocity = state[self._wheel_velocity_indices]\n",
    "\n",
    "        desired_wheel_velocity = self._vehicle_to_wheel_map @ command\n",
    "        torque = self._wheel_velocity_kp * (\n",
    "            desired_wheel_velocity - wheel_velocity\n",
    "        )\n",
    "\n",
    "        output.SetFromVector(torque)\n",
    "\n",
    "\n",
    "def teleop():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    time_step = 0.005\n",
    "    plant, scenegraph = AddMultibodyPlantSceneGraph(\n",
    "        builder, time_step=time_step\n",
    "    )\n",
    "    # plant.set_discrete_contact_solver(DiscreteContactSolver.kSap)\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    parser.AddModels(FindResource(\"models/mecanum/drake_obstacles.dmd.yaml\"))\n",
    "    plant.set_discrete_contact_solver(DiscreteContactSolver.kSap)\n",
    "    plant.Finalize()\n",
    "    robot_instance = plant.GetModelInstanceByName(\"mecanum_base\")\n",
    "\n",
    "    controller = builder.AddSystem(MyController(plant, robot_instance))\n",
    "    builder.Connect(\n",
    "        plant.get_state_output_port(), controller.get_input_port(1)\n",
    "    )\n",
    "    builder.Connect(\n",
    "        controller.get_output_port(), plant.get_actuation_input_port()\n",
    "    )\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.DeleteAddedControls()\n",
    "    meshcat.SetProperty(\"/Lights/PointLightNegativeX\", \"visible\", False)\n",
    "\n",
    "    gamepad = meshcat.GetGamepad()\n",
    "    if gamepad.index != None:  # use gamepad teleop\n",
    "        print(\"Gamepad found.  Using gamepad teleop.\")\n",
    "        gamepad_command = builder.AddSystem(GamepadCommand(meshcat))\n",
    "        builder.Connect(\n",
    "            gamepad_command.get_output_port(), controller.get_input_port(0)\n",
    "        )\n",
    "\n",
    "    else:  # use keyboard teleop\n",
    "        print(\"No gamepad found.  Using keyboard/slider teleop.\")\n",
    "        step = 0.02\n",
    "        meshcat.AddSlider(\n",
    "            name=\"xdot\",\n",
    "            increment_keycode=\"ArrowUp\",\n",
    "            decrement_keycode=\"ArrowDown\",\n",
    "            min=-3,\n",
    "            max=3,\n",
    "            step=step,\n",
    "            value=0,\n",
    "        )\n",
    "        meshcat.AddSlider(\n",
    "            name=\"ydot\",\n",
    "            increment_keycode=\"ArrowLeft\",\n",
    "            decrement_keycode=\"ArrowRight\",\n",
    "            min=-1,\n",
    "            max=1,\n",
    "            step=step,\n",
    "            value=0,\n",
    "        )\n",
    "        meshcat.AddSlider(\n",
    "            name=\"yawdot\",\n",
    "            increment_keycode=\"KeyD\",\n",
    "            decrement_keycode=\"KeyA\",\n",
    "            min=-1,\n",
    "            max=1,\n",
    "            step=step,\n",
    "            value=0,\n",
    "        )\n",
    "        sliders = builder.AddSystem(\n",
    "            MeshcatSliders(meshcat, [[\"xdot\", \"ydot\", \"yawdot\"]])\n",
    "        )\n",
    "        builder.Connect(\n",
    "            sliders.get_output_port(), controller.get_input_port(0)\n",
    "        )\n",
    "\n",
    "    # Add a meshcat visualizer.\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "        builder, scenegraph.get_query_output_port(), meshcat\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "    simulator.get_mutable_context()\n",
    "\n",
    "    if running_as_notebook:\n",
    "        simulator.set_target_realtime_rate(1.0)\n",
    "\n",
    "        meshcat.AddButton(\"Stop Simulation\", \"Escape\")\n",
    "        while meshcat.GetButtonClicks(\"Stop Simulation\") < 1:\n",
    "            simulator.AdvanceTo(simulator.get_context().get_time() + 2.0)\n",
    "        meshcat.DeleteButton(\"Stop Simulation\")\n",
    "    else:\n",
    "        simulator.AdvanceTo(0.1)\n",
    "\n",
    "\n",
    "teleop()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9e08339",
   "metadata": {
    "cell_id": "324c4ed0f4a84c1f9f7c9c58a241bc19",
    "deepnote_cell_type": "markdown"
   },
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "89815aca",
   "metadata": {
    "created_in_deepnote_cell": true,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "<a style='text-decoration:none;line-height:16px;display:flex;color:#5B5B62;padding:10px;justify-content:end;' href='https://deepnote.com?utm_source=created-in-deepnote-cell&projectId=57b741a1-558d-4f78-92d3-c473c39676c3' target=\"_blank\">\n",
    "<img alt='Created in deepnote.com' style='display:inline;max-height:16px;margin:0px;margin-right:7.5px;' src='' > </img>\n",
    "Created in <span style='font-weight:600;margin-left:4px;'>Deepnote</span></a>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
